{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Recommender Systems with DGL\n",
    "\n",
    "## Introduction\n",
    "\n",
    "Graph Neural Networks (GNN), as a methodology of learning representations on graphs, has gained much attention recently.  Various models such as Graph Convolutional Networks, GraphSAGE, etc. are proposed to obtain representations of whole graphs, or nodes on a single graph.\n",
    "\n",
    "A primary goal of Collaborative Filtering (CF) is to automatically make predictions about a user's interest, e.g. whether/how a user would interact with a set of items, given the interaction history of the user herself, as well as the histories of other users.  The user-item interaction can also be viewed as a bipartite graph, where users and items form two sets of nodes, and edges connecting them stands for interactions.  The problem can then be formulated as a *link-prediction* problem, where we try to predict whether an edge (of a given type) exists between two nodes.\n",
    "\n",
    "Based on this intuition, the academia developed multiple new models for CF, including but not limited to:\n",
    "\n",
    "* Geometric Learning Approaches\n",
    "  * [Geometric Matrix Completion](https://papers.nips.cc/paper/5938-collaborative-filtering-with-graph-information-consistency-and-scalable-methods.pdf)\n",
    "  * [Recurrent Multi-graph CNN](https://arxiv.org/pdf/1704.06803.pdf)\n",
    "* Graph-convolutional Approaches\n",
    "  * Models such as [R-GCN](https://arxiv.org/pdf/1703.06103.pdf) or [GraphSAGE](https://github.com/stellargraph/stellargraph/tree/develop/demos/link-prediction/hinsage) also apply.\n",
    "  * [Graph Convolutional Matrix Completion](https://arxiv.org/abs/1706.02263)\n",
    "  * [PinSage](https://arxiv.org/pdf/1806.01973.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dependencies\n",
    "\n",
    "* Latest DGL release: `conda install -c dglteam dgl`\n",
    "* `pandas`\n",
    "* `stanfordnlp`\n",
    "* `pytorch`\n",
    "* `tqdm` for displaying the progress bar."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading data\n",
    "\n",
    "In this tutorial, we focus on rating prediction on MovieLens-1M dataset.  The data comes from [MovieLens](http://files.grouplens.org/datasets/movielens/ml-1m.zip) and is shipped with the notebook already.\n",
    "\n",
    "After loading and train-validation-test-splitting the dataset, we process the movie title into (padded) word-ID sequences, and other features into categorical variables (i.e. integers).  We then store them as node features on the graph.\n",
    "\n",
    "Since user features and item features are different, we pad both types of features with zeros.\n",
    "\n",
    "All of the above is encapsulated in `movielens.MovieLens` class for clarity of this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import dgl\n",
    "import dgl.function as FN\n",
    "\n",
    "# Load Pytorch as backend\n",
    "dgl.load_backend('pytorch')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import movielens_torch as movielens\n",
    "import stanfordnlp\n",
    "\n",
    "# If you don't have stanfordnlp installed and the English models downloaded, please uncomment this statement\n",
    "#stanfordnlp.download('en', force=True)\n",
    "\n",
    "ml = movielens.MovieLens('ml-100k')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model\n",
    "\n",
    "We can now write a GraphSAGE layer.  In GraphSAGE, the node representation is updated with the representation in the previous layer as well as an aggregation (often mean) of \"messages\" sent from all neighboring nodes.\n",
    "\n",
    "### Algorithm\n",
    "\n",
    "The algorithm of a single GraphSAGE layer goes as follows for each node $v$:\n",
    "\n",
    "1. $h_{\\mathcal{N}(v)} \\gets \\mathtt{Average}_{u \\in \\mathcal{N}(v)} h_{u}$\n",
    "2. $h_{v} \\gets \\sigma\\left(W \\cdot \\mathtt{CONCAT}(h_v, h_{\\mathcal{N}(v)})\\right)$\n",
    "3. $h_{v} \\gets h_{v} / \\lVert h_{v} \\rVert_2$\n",
    "\n",
    "where\n",
    "\n",
    "* $\\mathtt{Average}$ can be replaced by any kind of aggregation including `sum`, `max`, or even an LSTM.\n",
    "* $\\sigma$ is any non-linearity function (e.g. `LeakyReLU`)\n",
    "\n",
    "We simply repeat the computation above for multiple GraphSAGE layers.\n",
    "\n",
    "### DGL Message Passing\n",
    "\n",
    "DGL adopts the message-passing paradigm, or scatter-apply-gather paradigm, for feature computation on a graph.  It decomposes the computation into three stages:\n",
    "\n",
    "1. *Message computation*: each edge is computed a message according to features on the edge itself, as well as the features on its source and destination node.  Often times, the message computation simply involves copying the representation of the source node.\n",
    "2. *Message aggregation*: each node then \"receives\" the messages sent from its neighbors, and call a function which reduces these messages into a single representation independent of the number of neighbors.  Averaging and summing are two of the most common message aggregation functions.\n",
    "3. *Node feature update*: with an aggregated representation from the neighbors, a node then updates its own representation using the aggregation.\n",
    "\n",
    "With the three stages in mind, we can easily figure out how to map the GraphSAGE layer computation into the message-passing paradigm:\n",
    "\n",
    "1. $h_{\\mathcal{N}(v)} \\gets \\underbrace{\\mathtt{Average}_{u \\in \\mathcal{N}(v)} \\underbrace{h_{u}}_{\\text{Message computation (copy from source)}}}_{\\text{Message aggregation}}$\n",
    "2. $h_{v} \\gets \\underbrace{\\sigma\\left(W \\cdot \\mathtt{CONCAT}(h_v, h_{\\mathcal{N}(v)})\\right)}_{\\text{Node feature update}}$\n",
    "3. $h_{v} \\gets \\underbrace{h_{v} / \\lVert h_{v} \\rVert_2}_{\\text{Node feature update}}$\n",
    "\n",
    "While DGL does not provide the $\\mathtt{Average}$ aggregation function yet (as it's a future work item), it does provide the $\\mathtt{Sum}$ aggregation.  So we can modify the algorithm above to the following that is readily to be implemented in DGL:\n",
    "\n",
    "1. $d_{\\mathcal{N}(v)} \\gets \\underbrace{\\mathtt{Sum}_{u \\in \\mathcal{N}(v)} \\underbrace{1}_{\\text{Message computation (copy from source)}}}_{\\text{Message aggregation}}$\n",
    "2. $h_{\\mathcal{N}(v)} \\gets \\underbrace{\\mathtt{Sum}_{u \\in \\mathcal{N}(v)} \\underbrace{h_{u}}_{\\text{Message computation (copy from source)}}}_{\\text{Message aggregation}}$\n",
    "3. $h_{v} \\gets \\underbrace{\\sigma\\left(W \\cdot \\mathtt{CONCAT}(h_v, h_{\\mathcal{N}(v)} / d_{\\mathcal{N}(v)})\\right)}_{\\text{Node feature update}}$\n",
    "4. $h_{v} \\gets \\underbrace{h_{v} / \\lVert h_{v} \\rVert_2}_{\\text{Node feature update}}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sampling\n",
    "\n",
    "Ideally, we wish to execute a full update of the node embeddings with the GraphSAGE layer.  However, when the graph scales up, the full update soon becomes impractical, because the node embeddings couldn't fit in the GPU memory.\n",
    "\n",
    "A natural solution would be partitioning the nodes and computing the embeddings one partition (minibatch) at a time.  The nodes at one convolution layer then only depends on their neighbors, rather than all the nodes in the graph, hence reducing the computational cost.  However, if we have multiple layers, and some of the nodes have a lot of neighbors (which is often the case since the degree distribution of many real-world graphs follow [power-law](https://en.wikipedia.org/wiki/Scale-free_network)), then the computation may still eventually depend on every node in the graph.\n",
    "\n",
    "*Neighbor sampling* is an answer to further reduce the cost of computing node embeddings.  When aggregating messages, instead of collecting from all neighboring nodes, we only collect from some of the randomly-sampled (for instance, uniform sampling at most K neighbors without replacement) neighbors.\n",
    "\n",
    "DGL provides the `NodeFlow` object that describes the computation dependency of nodes in a graph convolutional network, as well as various samplers that constructs such `NodeFlow`s as graphs.  From a programmer's perspective, training with minibatch and neighbor sampling reduces to propagating the messages in `NodeFlow` as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mix_embeddings(ndata, emb, proj):\n",
    "    \"\"\"Adds external (categorical and numeric) features into node representation G.ndata['h']\"\"\"\n",
    "    extra_repr = []\n",
    "    for key, value in ndata.items():\n",
    "        if (value.dtype == torch.int64) and key in emb:\n",
    "            result = emb[key](value)\n",
    "            if result.dim() == 3:    # bag of words: the result would be a (n_nodes x seq_len x feature_size) tensor\n",
    "                result = result.mean(1)\n",
    "            extra_repr.append(result)\n",
    "        elif (value.dtype == torch.float32) and key in proj:\n",
    "            result = proj[key](value)\n",
    "            extra_repr.append(result)\n",
    "    ndata['h'] = ndata['h'] + torch.stack(extra_repr, 0).sum(0)\n",
    "    \n",
    "def init_weight(param, initializer, nonlinearity):\n",
    "    initializer = getattr(nn.init, initializer)\n",
    "    if nonlinearity is not None:\n",
    "        initializer(param)\n",
    "    else:\n",
    "        initializer(param, nn.init.calculate_gain(nonlinearity))\n",
    "        \n",
    "def init_bias(param):\n",
    "    nn.init.constant_(param, 0)\n",
    "\n",
    "class GraphSageConvWithSampling(nn.Module):\n",
    "    def __init__(self, feature_size):\n",
    "        super(GraphSageConvWithSampling, self).__init__()\n",
    "\n",
    "        self.feature_size = feature_size\n",
    "        self.W = nn.Linear(feature_size * 2, feature_size)\n",
    "        init_weight(self.W.weight, 'xavier_uniform_', 'leaky_relu')\n",
    "        init_bias(self.W.bias)\n",
    "\n",
    "    def forward(self, nodes):\n",
    "        h_agg = nodes.data['h_agg']\n",
    "        h = nodes.data['h']\n",
    "        w = nodes.data['w'][:, None]\n",
    "        h_agg = (h_agg-h)/(w-1).clamp(min=1)    # HACK 1\n",
    "        h_concat = torch.cat([h, h_agg], 1)\n",
    "        h_new = F.leaky_relu(self.W(h_concat))\n",
    "        return {'h': h_new / h_new.norm(dim=1, keepdim=True).clamp(min=1e-6)}\n",
    "    \n",
    "class GraphSageWithSampling(nn.Module):\n",
    "    def __init__(self, feature_size, n_layers, G):\n",
    "        super(GraphSageWithSampling, self).__init__()\n",
    "        \n",
    "        self.feature_size = feature_size\n",
    "        self.n_layers = n_layers\n",
    "\n",
    "        self.convs = nn.ModuleList([GraphSageConvWithSampling(feature_size) for _ in range(n_layers)])\n",
    "        \n",
    "        self.emb = nn.ModuleDict()\n",
    "        self.proj = nn.ModuleDict()\n",
    "\n",
    "        for key, scheme in G.node_attr_schemes().items():\n",
    "            if scheme.dtype == torch.int64:\n",
    "                n_items = G.ndata[key].max().item()\n",
    "                self.emb[key] = nn.Embedding(\n",
    "                        n_items + 1,\n",
    "                        self.feature_size,\n",
    "                        padding_idx=0)\n",
    "                nn.init.normal_(self.emb[key].weight, 1 / self.feature_size)\n",
    "            elif scheme.dtype == torch.float32:\n",
    "                w = nn.Linear(scheme.shape[0], self.feature_size)\n",
    "                init_weight(w.weight, 'xavier_uniform_', 'leaky_relu')\n",
    "                init_bias(w.bias)\n",
    "                self.proj[key] = nn.Sequential(w, nn.LeakyReLU())\n",
    "                \n",
    "        self.G = G\n",
    "        \n",
    "        self.node_emb = nn.Embedding(G.number_of_nodes() + 1, feature_size)\n",
    "        nn.init.normal_(self.node_emb.weight, std=1 / self.feature_size)\n",
    "\n",
    "    msg = [FN.copy_src('h', 'h'),\n",
    "           FN.copy_src('one', 'one')]\n",
    "    red = [FN.sum('h', 'h_agg'), FN.sum('one', 'w')]\n",
    "\n",
    "    def forward(self, nf):\n",
    "        '''\n",
    "        nf: NodeFlow.\n",
    "        '''\n",
    "        nf.copy_from_parent(edge_embed_names=None)\n",
    "        for i in range(nf.num_layers):\n",
    "            nf.layers[i].data['h'] = self.node_emb(nf.layer_parent_nid(i) + 1)\n",
    "            nf.layers[i].data['one'] = torch.ones(nf.layer_size(i))\n",
    "            mix_embeddings(nf.layers[i].data, model.gcn.emb, model.gcn.proj)\n",
    "        if self.n_layers == 0:\n",
    "            return nf.layers[i].data['h']\n",
    "        for i in range(self.n_layers):\n",
    "            nf.block_compute(i, self.msg, self.red, self.convs[i])\n",
    "\n",
    "        result = nf.layers[self.n_layers].data['h']\n",
    "        assert (result != result).sum() == 0\n",
    "        return result\n",
    "    \n",
    "class GraphSAGERecommender(nn.Module):\n",
    "    def __init__(self, gcn):\n",
    "        super(GraphSAGERecommender, self).__init__()\n",
    "        \n",
    "        self.gcn = gcn\n",
    "        self.node_biases = nn.Parameter(torch.zeros(gcn.G.number_of_nodes()+1))\n",
    "        \n",
    "    def forward(self, nf, src, dst):\n",
    "        h_output = self.gcn(nf)\n",
    "        h_src = h_output[nodeflow.map_from_parent_nid(-1, src, True)]\n",
    "        h_dst = h_output[nodeflow.map_from_parent_nid(-1, dst, True)]\n",
    "        score = (h_src * h_dst).sum(1) + self.node_biases[src+1] + self.node_biases[dst+1]\n",
    "        return score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "As above, training now only involves\n",
    "1. Initializing a sampler\n",
    "2. Iterating over the neighbor sampler, propagating the messages, and computing losses and gradients as usual.\n",
    "\n",
    "Meanwhile, we also evaluate the RMSE on validation and test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import tqdm\n",
    "import spotlight\n",
    "import pickle\n",
    "\n",
    "g = ml.g\n",
    "# Find the subgraph of all \"training\" edges\n",
    "g_train = g.edge_subgraph(g.filter_edges(lambda edges: edges.data['train']), True)\n",
    "g_train.copy_from_parent()\n",
    "g_train.readonly()\n",
    "eid_valid = g.filter_edges(lambda edges: edges.data['valid'])\n",
    "eid_test = g.filter_edges(lambda edges: edges.data['test'])\n",
    "src_valid, dst_valid = g.find_edges(eid_valid)\n",
    "src_test, dst_test = g.find_edges(eid_test)\n",
    "src, dst = g_train.all_edges()\n",
    "rating = g_train.edata['rating']\n",
    "rating_valid = g.edges[eid_valid].data['rating']\n",
    "rating_test = g.edges[eid_test].data['rating']\n",
    "\n",
    "model = GraphSAGERecommender(GraphSageWithSampling(100, 1, g_train))\n",
    "opt = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-9)\n",
    "\n",
    "batch_size = 1024\n",
    "n_users = len(ml.user_ids)\n",
    "n_products = len(ml.product_ids)\n",
    "\n",
    "for epoch in range(50):\n",
    "    model.eval()\n",
    "    \n",
    "    # Validation & Test, we precompute GraphSage output for all nodes first.\n",
    "    sampler = dgl.contrib.sampling.NeighborSampler(\n",
    "        g_train,\n",
    "        batch_size,\n",
    "        5,\n",
    "        1,\n",
    "        seed_nodes=torch.arange(g.number_of_nodes()),\n",
    "        prefetch=True,\n",
    "        add_self_loop=True,\n",
    "        shuffle=False,\n",
    "        num_workers=4\n",
    "    )\n",
    "\n",
    "    with torch.no_grad():\n",
    "        h = []\n",
    "        for nf in sampler:\n",
    "            #import pdb\n",
    "            #pdb.set_trace()\n",
    "            h.append(model.gcn.forward(nf))\n",
    "        h = torch.cat(h)\n",
    "\n",
    "        # Compute validation RMSE\n",
    "        score = torch.zeros(len(src_valid))\n",
    "        for i in range(0, len(src_valid), batch_size):\n",
    "            s = src_valid[i:i+batch_size]\n",
    "            d = dst_valid[i:i+batch_size]\n",
    "            score[i:i+batch_size] = (h[s] * h[d]).sum(1) + model.node_biases[s + 1] + model.node_biases[d + 1]\n",
    "        valid_rmse = ((score - rating_valid) ** 2).mean().sqrt()\n",
    "\n",
    "        # Compute test RMSE\n",
    "        score = torch.zeros(len(src_test))\n",
    "        for i in range(0, len(src_test), batch_size):\n",
    "            s = src_test[i:i+batch_size]\n",
    "            d = dst_test[i:i+batch_size]\n",
    "            score[i:i+batch_size] = (h[s] * h[d]).sum(1) + model.node_biases[s + 1] + model.node_biases[d + 1]\n",
    "        test_rmse = ((score - rating_test) ** 2).mean().sqrt()\n",
    "        \n",
    "    model.train()\n",
    "    \n",
    "    shuffle_idx = torch.randperm(g_train.number_of_edges())\n",
    "    src_shuffled = src[shuffle_idx]\n",
    "    dst_shuffled = dst[shuffle_idx]\n",
    "    rating_shuffled = rating[shuffle_idx]\n",
    "    src_batches = src_shuffled.split(batch_size)\n",
    "    dst_batches = dst_shuffled.split(batch_size)\n",
    "    rating_batches = rating_shuffled.split(batch_size)\n",
    "\n",
    "    seed_nodes = torch.cat(sum([[s, d] for s, d in zip(src_batches, dst_batches)], []))\n",
    "    \n",
    "    sampler = dgl.contrib.sampling.NeighborSampler(\n",
    "        g_train,               # the graph\n",
    "        batch_size * 2,        # number of nodes to compute at a time, HACK 2\n",
    "        5,                     # number of neighbors for each node\n",
    "        1,                     # number of layers in GCN\n",
    "        seed_nodes=seed_nodes, # list of seed nodes, HACK 2\n",
    "        prefetch=True,         # whether to prefetch the NodeFlows\n",
    "        add_self_loop=True,    # whether to add a self-loop in the NodeFlows, HACK 1\n",
    "        shuffle=False,         # whether to shuffle the seed nodes.  Should be False here.\n",
    "        num_workers=4,\n",
    "    )\n",
    "\n",
    "    # Training\n",
    "    for s, d, r, nodeflow in zip(src_batches, dst_batches, rating_batches, sampler):\n",
    "        score = model.forward(nodeflow, s, d)\n",
    "        loss = ((score - r) ** 2).mean()\n",
    "        \n",
    "        opt.zero_grad()\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "\n",
    "    print('Training loss:', loss.item(), 'Validation RMSE:', valid_rmse.item(), 'Test RMSE:', test_rmse.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
